-
Notifications
You must be signed in to change notification settings - Fork 293
Enable range learning for QAT #2033
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/2033
Note: Links to docs will display an error until the docs builds have been completed. ✅ No FailuresAs of commit 045b99e with merge base 0aa8dbd ( This comment was automatically generated by Dr. CI and updates every 15 minutes. |
@andrewor14 has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator. |
cf3e00e
to
f3903d4
Compare
@andrewor14 has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator. |
f3903d4
to
3277221
Compare
@andrewor14 has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator. |
3277221
to
b33281e
Compare
@andrewor14 has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator. |
b33281e
to
ec45945
Compare
@andrewor14 has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator. |
912a925
to
cb98579
Compare
I think this needs a rebase |
cb98579
to
fbd4716
Compare
@andrewor14 has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator. |
**Summary:** This commit adds the option for QAT users to use range learning during training. Range learning means we train the scale and zero point instead of recomputing them based on the input at every iteration. Example usage: ``` import torch from torchao.quantization import quantize_ from torchao.quantization.qat import ( FakeQuantizeConfig, IntXQuantizationAwareTrainingConfig, initialize_fake_quantizers, ) config = FakeQuantizeConfig( torch.int8, "per_channel", is_dynamic=False, range_learning=True, scale_precision=torch.float32, zero_point_precision=torch.float32, ) m = M() example_inputs = (torch.randn(16, 32),) quantize_(m, IntXQuantizationAwareTrainingConfig(weight_config=config)) # New required step to turn scales and zero points into trainable # `nn.Parameters`, must be called before initializing the optimizer initialize_fake_quantizers(m, example_inputs) # initialize the optimizer # do training ``` **Test Plan:** python test/quantization/test_qat.py -k test_fake_quantize_config_dynamic_and_range_learning python test/quantization/test_qat.py -k test_fake_quantizer_range_learning python test/quantization/test_qat.py -k test_qat_range_learning
fbd4716
to
045b99e
Compare
@andrewor14 has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator. |
@andrewor14 Did you downgrade cutlass to 3.8 by mistake in this PR? I'm trying to build from source and it fails to build the MX kernels. |
Ah my mistake, let me look into a fix |
Cutlass was accidentally downgraded to 3.8 in #2033. This commit undoes this change and brings cutlass back to the original 3.9 commit.
Summary: This commit adds the option for QAT users to use range learning during training. Range learning means we train the scale and zero point instead of recomputing them based on the input in every iteration.
Example usage:
Test Plan:
python test/quantization/test_qat.py -k test_fake_quantize_config_dynamic_and_range_learning
python test/quantization/test_qat.py -k test_fake_quantizer_range_learning
python test/quantization/test_qat.py -k test_qat_range_learning